#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2024 Song Yan 843558935@qq.com. All rights reserved.
# This code is licensed under the GPL-2.0+ license. See LICENSE file in the project root for full license information.

import logging
import os
from invoke import task

from tasks.constants import DNF_CONF_PATH, DNF_UPGRADE_DATA, RELEASE_VERSION
from tasks.utils import manage_rpmdb_symlinks, replace_dbpath, safe_run_command
from tasks.record_status import status_manager


logger = logging.getLogger(__name__)

@task
@status_manager.skip_if_success_decorator
def rebuild_rpmdb(ctx):
    logger.info("Rebuilding RPM database...")
    rpm_db_path = "/usr/lib/sysimage/rpm"
    

    # 移动或备份旧的rpm数据库路径
    if not os.path.isdir(rpm_db_path):
        os.rename("/var/lib/rpm", rpm_db_path)

    # 创建迁移标记文件和符号链接
    manage_rpmdb_symlinks()

    # 更新宏文件中的数据库路径
    macros_file_path = "/usr/lib/rpm/macros"
    
    with open(macros_file_path, "r+") as macros_file:
        content = macros_file.read()
        updated_content = replace_dbpath(content)
        macros_file.seek(0)
        macros_file.write(updated_content)
        macros_file.truncate()

    safe_run_command(ctx, "rpm --rebuilddb")

    return True

@task
@status_manager.skip_if_success_decorator
def remove_modules(ctx):
    logger.info("Removing modules...")
    command = "rm -rf /etc/dnf/modules.d/*.module"
    if safe_run_command(ctx, command).ok:
        logger.info("Modules removed successfully")
    else:
        logger.error("Failed to remove modules")
        return False
    return True

@task
@status_manager.skip_if_success_decorator
def fix_release_version(ctx):
    logger.info("Fixing release version...")
    command = "if rpm -q opencloudos-release; then rpm -e --nodeps opencloudos-release; fi"
    if safe_run_command(ctx, command).ok:
        logger.info("Release version fixed successfully")
    else:
        logger.error("Failed to fix release version")
        return False
    return True

@task
@status_manager.skip_if_success_decorator
def setup_repositories(ctx):
    logger.info("Setting up repositories for opencloudos {}...".format(RELEASE_VERSION))
    
    with open(DNF_CONF_PATH, 'w') as f:
        f.write("""
[main]
gpgcheck=0
max_parallel_downloads=10
installonly_limit=3
clean_requirements_on_remove=True
best=False
skip_if_unavailable=True
deltarpm=False
zchunk=False
allowerasing=True

[BaseOS-oc9]
name=OpenCloudOS $releasever - BaseOS-oc9
baseurl=https://mirrors.opencloudos.tech/opencloudos/$releasever/BaseOS/$basearch/os/
gpgcheck=0
enabled=1

[AppStream-oc9]
name=OpenCloudOS $releasever - AppStream-oc9
baseurl=https://mirrors.opencloudos.tech/opencloudos/$releasever/AppStream/$basearch/os/
gpgcheck=0
enabled=1
""")
    
    command = f"dnf -c {DNF_CONF_PATH} makecache -y --disablerepo='*' --enablerepo='*-oc9'"

    if safe_run_command(ctx, command).ok:
        logger.info("DNF makecache completed successfully")
    else:
        logger.error("DNF makecache failed")
        return False
    return True

@task(pre=[setup_repositories, rebuild_rpmdb, remove_modules])
@status_manager.skip_if_success_decorator
def prepare_upgrade(ctx):
    logger.info("Preparing upgrade...")
    to_distro_sync = ["distro-sync"]
    to_remove= [
                "remove NetworkManager-bluetooth",
                "remove NetworkManager-team",
                "remove NetworkManager-wifi",
                "remove NetworkManager-wwan",
                ]

    to_swap = ["swap mozjs60 mozjs102"]
    to_install = []
    
    command = f"dnf -c {DNF_CONF_PATH} distro-sync -y --releasever={RELEASE_VERSION} --allowerasing --nobest --disablerepo='*' --enablerepo='*-oc9'"
    result = safe_run_command(ctx, command)
    
    for line in result.stderr.splitlines():
        if "conflicts" not in line:
            continue
        parts = line.strip().split()
        
        xline = "swap {} {}".format(parts[-1], parts[-7])
        if xline not in to_swap:
            to_swap.append(xline)

    if safe_run_command(ctx, "rpm -q --quiet dejavu-sans-fonts").ok:
        to_swap.append("swap dejavu-*-fonts dejavu-fonts")
        
    if os.path.isdir("/sys/firmware/efi/"):
        to_install.append("install shim")
    
    with open (DNF_UPGRADE_DATA, "w") as f:
        f.write("\n".join(to_distro_sync + to_remove + to_swap + to_install))
    logger.info("Upgrade prepared successfully")

    return True
